import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import os
import pandas as pd

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def train_and_visualize(cases,all_data, test_size=0.2, random_state=42, output_dir="results"):
    """
    训练随机森林回归模型并可视化结果
    
    参数:
        all_data: 包含6个DataFrame的列表，每个DataFrame包含'x', 't', 'u'三列
        test_size: 测试集比例，默认为0.2
        random_state: 随机种子，确保结果可复现
        output_dir: 结果保存目录
    """
    # 创建保存结果的目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 对每个DataFrame训练一个模型
    for i, df in enumerate(all_data, 1):
        print(f"\n处理第 {i} 个数据集...")
        
        # 准备特征和目标变量
        X = df[['x', 't']]
        y = df['u']
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=random_state
        )
        
        # 训练随机森林模型
        model = RandomForestRegressor(
            n_estimators=100, 
            max_depth=None,
            min_samples_split=2,
            random_state=random_state,
            n_jobs=-1  # 使用所有CPU核心
        )
        model.fit(X_train, y_train)
        
        # 在测试集上进行预测
        y_pred = model.predict(X_test)
        
        # 计算误差
        mse = mean_squared_error(y_test, y_pred)
        print(f"模型 {i} 测试集MSE: {mse:.4f}")
        
        y_pred = model.predict(X)
        # 可视化结果
        fig, axes = plt.subplots(1, 2, figsize=(16, 6))
        fig.suptitle(f" {cases[i-1]}", fontsize=16)
        
        # 1. 预测结果热力图
        scatter = axes[0].scatter(
            X['x'], X['t'], 
            c=y_pred, 
            cmap='jet', 
            s=1,
            alpha=1
        )
        axes[0].set_title('测试集预测结果热力图')
        axes[0].set_xlabel('x')
        axes[0].set_ylabel('t')
        fig.colorbar(scatter, ax=axes[0], label='预测值')
        

        # 2. 误差热力图
        error = y - y_pred
        scatter_error = axes[1].scatter(
            X['x'], X['t'], 
            c=error, 
            cmap='jet', 
            s=1,
            alpha=1
        )
        axes[1].set_title('测试集预测误差热力图')
        axes[1].set_xlabel('x')
        axes[1].set_ylabel('t')
        fig.colorbar(scatter_error, ax=axes[1], label='误差值')
        
        plt.tight_layout(rect=[0, 0, 1, 0.96])  # 调整布局，为suptitle留出空间
        plt.show()
        
    
    print("\n所有模型训练和可视化完成！")